#include <crypto/skcipher.h>
#include <linux/crypto.h>
#include <linux/err.h>
#include <linux/init.h>
#include <linux/kernel.h>
#include <linux/mm.h>
#include <linux/module.h>
#include <linux/scatterlist.h>
#include <linux/slab.h>
#include <linux/string.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Kernel Crypto Example");
MODULE_DESCRIPTION("Linux内核6.6版本下的AES对称加密示例（修正版）");

// 原始明文数据
static char *plaintext = "Hello, 这是Linux内核加密测试消息!";
static char key[32] = "12345678901234567890123456789012"; // 256位密钥
static char iv[16] = "1234567890123456";                  // 128位初始化向量
static char iv2[16] = "1234567890123456";                 // 128位初始化向量

static struct crypto_skcipher *tfm;  // 加密算法句柄
static struct skcipher_request *req; // 加密请求结构
static char *input_data;             // 输入数据缓冲区
static char *output_data;            // 输出数据缓冲区
static char *decrypt_data;           // 解密数据缓冲区

static int __init crypto_init(void) {
  int ret = 0;
  struct scatterlist sg_input, sg_output, sg_decrypt;
  size_t pt_len = strlen(plaintext);
  size_t padded_len = (pt_len + 15) & ~15; // 向上舍入到16的倍数

  printk(KERN_INFO "正在初始化Linux内核6.6加密模块（修正版）\n");

  // 分配对称加密算法句柄，使用AES-CBC模式
  tfm = crypto_alloc_skcipher("cbc(aes)", 0, CRYPTO_ALG_ASYNC);
  if (IS_ERR(tfm)) {
    printk(KERN_ERR "分配skcipher句柄失败\n");
    return PTR_ERR(tfm);
  }

  // 设置加密密钥
  ret = crypto_skcipher_setkey(tfm, key, 32);
  if (ret) {
    printk(KERN_ERR "设置密钥失败\n");
    crypto_free_skcipher(tfm);
    return ret;
  }

  // 分配加密请求结构
  req = skcipher_request_alloc(tfm, GFP_KERNEL);
  if (!req) {
    printk(KERN_ERR "分配请求结构失败\n");
    crypto_free_skcipher(tfm);
    return -ENOMEM;
  }

  // 使用__get_free_page分配连续的物理内存页（可以用alloc_pages、kmalloc、kzalloc代替）
  input_data = (char *)__get_free_page(GFP_KERNEL);
  output_data = (char *)__get_free_page(GFP_KERNEL);
  decrypt_data = (char *)__get_free_page(GFP_KERNEL);

  if (!input_data || !output_data || !decrypt_data) {
    printk(KERN_ERR "分配连续物理内存失败\n");
    ret = -ENOMEM;
    goto cleanup_pages;
  }

  // 复制明文到输入缓冲区并用零填充
  memset(input_data, 0, PAGE_SIZE);
  memset(output_data, 0, PAGE_SIZE);
  memset(decrypt_data, 0, PAGE_SIZE);
  strcpy(input_data, plaintext);

  // 设置加密操作的scatterlist
  sg_init_one(&sg_input, input_data, padded_len);
  sg_init_one(&sg_output, output_data, padded_len);

  // 设置加密请求参数
  skcipher_request_set_crypt(req, &sg_input, &sg_output, padded_len, iv);

  // 执行加密操作
  ret = crypto_skcipher_encrypt(req);
  if (ret) {
    printk(KERN_ERR "加密操作失败: %d\n", ret);
    goto cleanup_pages;
  }

  printk(KERN_INFO "加密操作成功\n");
  printk(KERN_INFO "原始明文: %s\n", plaintext);
  printk(KERN_CONT "加密后密文(十六进制): ");
  for (int i = 0; i < padded_len; i++) {
    printk(KERN_CONT "%02x", (unsigned char)output_data[i]);
  }
  printk(KERN_CONT "\n");

  // 现在执行解密操作以验证加密结果
  sg_init_one(&sg_input, output_data, padded_len);
  sg_init_one(&sg_decrypt, decrypt_data, padded_len);

  skcipher_request_set_crypt(req, &sg_input, &sg_decrypt, padded_len, iv2);

  ret = crypto_skcipher_decrypt(req);
  if (ret) {
    printk(KERN_ERR "解密操作失败: %d\n", ret);
    goto cleanup_pages;
  }

  printk(KERN_INFO "解密操作成功\n");
  printk(KERN_INFO "解密后文本: %s\n", decrypt_data);

cleanup_pages:
  // 释放分配的连续物理内存页
  if (input_data)
    free_page((unsigned long)input_data);
  if (output_data)
    free_page((unsigned long)output_data);
  if (decrypt_data)
    free_page((unsigned long)decrypt_data);

  skcipher_request_free(req);
  crypto_free_skcipher(tfm);

  return ret;
}

static void __exit crypto_exit(void) {
  printk(KERN_INFO "正在退出Linux内核6.6加密模块（修正版）\n");
}

module_init(crypto_init);
module_exit(crypto_exit);